package com.star.spring.ratelimit.limiter;

import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * 本地令牌桶限流器
 * 
 * @author star
 * @since 2025-01-24
 */
@Slf4j
@Component
public class LocalTokenBucketLimiter {


    /**
     * 存储不同key的自定义令牌桶
     */
    private final ConcurrentHashMap<String, TokenBucket> tokenBuckets = new ConcurrentHashMap<>();


    /**
     * 尝试获取令牌（根据限流配置）
     * 
     * @param key 限流key
     * @param maxRequests 最大请求数
     * @param timeWindowSeconds 时间窗口（秒）
     * @return 是否获取成功
     */
    public boolean tryAcquire(String key, int maxRequests, int timeWindowSeconds) {
        // 计算桶容量：允许突发流量，桶容量设为时间窗口内最大请求数
        int bucketCapacity = maxRequests;
        
        // 获取或创建自定义令牌桶
        TokenBucket tokenBucket = tokenBuckets.computeIfAbsent(key, k -> {
            log.info("Creating custom token bucket - key: {}, maxRequests: {}, bucketCapacity: {}, timeWindow: {}s", 
                    key, maxRequests, bucketCapacity, timeWindowSeconds);
            return new TokenBucket(bucketCapacity, maxRequests, timeWindowSeconds);
        });

        boolean acquired = tokenBucket.tryAcquire(1);
        log.debug("Custom token bucket limiter - key: {}, maxRequests: {}, timeWindow: {}s, bucketCapacity: {}, currentTokens: {}, acquired: {}", 
                key, maxRequests, timeWindowSeconds, bucketCapacity, tokenBucket.getCurrentTokens(), acquired);
        return acquired;
    }

}
